package payserver

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"encoding/base64"
	"net/http"

	"google.golang.org/grpc/status"
)

var _ Crypto = (*AesCBC)(nil)
var _ Crypto = (*AesECB)(nil)

type (
	AesCBC struct {
		key   []byte
		iv    []byte
		block cipher.Block
	}

	AesECB struct {
		key   []byte
		block cipher.Block
	}

	AesCBCOption func(*AesCBC)
	AesECBOption func(*AesECB)
)

const (
	defaultAesCBCKey = "759e7e0e85e2e927382b0e1259637d41"
	defaultAesCBCIv  = "bdc3ad6584ca3b4a"
	defaultAesECBKey = "759e7e0e85e2e927382b0e1259637d41"
)

// NewAesCBC 实例化Aes加密
func NewAesCBC(options ...AesCBCOption) (Crypto, error) {
	a := &AesCBC{
		key: []byte(defaultAesCBCKey),
		iv:  []byte(defaultAesCBCIv),
	}

	for _, option := range options {
		option(a)
	}

	var err error
	a.block, err = aes.NewCipher(a.key)
	if err != nil {
		return nil, err
	}

	return a, nil
}

// NewAesECB 实例化Aes加密
func NewAesECB(options ...AesECBOption) (Crypto, error) {
	a := &AesECB{
		key: []byte(defaultAesECBKey),
	}

	for _, option := range options {
		option(a)
	}

	var err error
	a.block, err = aes.NewCipher(a.key)
	if err != nil {
		return nil, err
	}

	return a, nil
}

func WithAesCBCKey(key string) AesCBCOption {
	return func(a *AesCBC) {
		a.key = []byte(key)
	}
}

func WithAesCBCIv(iv string) AesCBCOption {
	return func(a *AesCBC) {
		a.iv = []byte(iv)
	}
}

func WithAesECBKey(key string) AesECBOption {
	return func(a *AesECB) {
		a.key = []byte(key)
	}
}

func (a *AesCBC) Encrypt(data []byte) (s []byte, err error) {
	defer func() {
		if p := recover(); p != nil {
			err = status.Errorf(http.StatusInternalServerError, "AesCBC Decrypt: %v", p)
		}
	}()
	origData := pkCS5Padding(data, a.block.BlockSize())
	crypt := make([]byte, len(origData))
	bm := cipher.NewCBCEncrypter(a.block, a.iv)
	bm.CryptBlocks(crypt, origData)

	return []byte(base64.StdEncoding.EncodeToString(crypt)), err
}

func (a *AesCBC) Decrypt(data []byte) (s []byte, err error) {
	defer func() {
		if p := recover(); p != nil {
			err = status.Errorf(http.StatusInternalServerError, "AesCBC Decrypt: %v", p)
		}
	}()
	crypt, err := base64.StdEncoding.DecodeString(string(data))
	if err != nil {
		return nil, err
	}
	origData := make([]byte, len(crypt))
	bm := cipher.NewCBCDecrypter(a.block, a.iv)
	bm.CryptBlocks(origData, crypt)
	origData = pkCS5UnPadding(origData)
	return origData, err
}

func (a *AesECB) Encrypt(data []byte) (s []byte, err error) {
	defer func() {
		if p := recover(); p != nil {
			err = status.Errorf(http.StatusInternalServerError, "AesECB Encrypt: %v", p)
		}
	}()
	blockSize := a.block.BlockSize()
	plaintext := pkCS5Padding(data, blockSize)

	ciphertext := make([]byte, len(plaintext))
	for bs, be := 0, blockSize; bs < len(plaintext); bs, be = bs+blockSize, be+blockSize {
		a.block.Encrypt(ciphertext[bs:be], plaintext[bs:be])
	}

	return []byte(base64.StdEncoding.EncodeToString(ciphertext)), err
}

func (a *AesECB) Decrypt(data []byte) (s []byte, err error) {
	defer func() {
		if p := recover(); p != nil {
			err = status.Errorf(http.StatusInternalServerError, "AesECB Decrypt: %v", p)
		}
	}()
	var crypt []byte
	crypt, err = base64.StdEncoding.DecodeString(string(data))
	if err != nil {
		return nil, err
	}
	ciphertext := crypt

	blockSize := a.block.BlockSize()
	plaintext := make([]byte, len(ciphertext))

	for bs, be := 0, blockSize; bs < len(ciphertext); bs, be = bs+blockSize, be+blockSize {
		a.block.Decrypt(plaintext[bs:be], ciphertext[bs:be])
	}

	return pkCS5UnPadding(plaintext), err
}

func pkCS5Padding(ciphertext []byte, blockSize int) []byte {
	padding := blockSize - len(ciphertext)%blockSize
	padText := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(ciphertext, padText...)
}

func pkCS5UnPadding(origData []byte) []byte {
	length := len(origData)
	unPadding := int(origData[length-1])
	return origData[:(length - unPadding)]
}
